package cn.iocoder.yudao.framework.server.websocket.security;

import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.iocoder.yudao.framework.cache.NettyChannelCacheService;
import cn.iocoder.yudao.framework.cache.NettyWsUserChannelCacheService;
import cn.iocoder.yudao.framework.common.enums.UserTypeEnum;
import cn.iocoder.yudao.framework.handlermapping.NettyWsChannelContext;
import cn.iocoder.yudao.framework.handlermapping.WsSessionContext;
import cn.iocoder.yudao.framework.prometheus.ws.WebSocketConnectionMetrics;
import cn.iocoder.yudao.framework.server.websocket.WebSocketAuthSuccessEvent;
import cn.iocoder.yudao.module.system.api.oauth2.OAuth2TokenApi;
import cn.iocoder.yudao.module.system.api.oauth2.dto.OAuth2AccessTokenCheckRespDTO;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.handler.codec.http.*;
import io.netty.util.Attribute;
import io.netty.util.AttributeKey;
import lombok.extern.slf4j.Slf4j;

import javax.annotation.Resource;
import java.util.Objects;

@Slf4j
public class WebTokenAuthorizationHandler extends ChannelInboundHandlerAdapter {


    @Resource
    private OAuth2TokenApi oauth2TokenApi;
    private WebSocketUser webSocketUser;

    @Resource
    private WebSocketConnectionMetrics webSocketConnectionMetrics;


    private WebSocketUser checkAccessToken(String token) {
        token = SecurityFrameworkUtils.obtainAuthorization(token);
        if (StrUtil.isNotEmpty(token)) {
            OAuth2AccessTokenCheckRespDTO accessToken = oauth2TokenApi.checkAccessToken(token);
            if (accessToken == null) {
                return null;
            }

            final UserTypeEnum userType = UserTypeEnum.valueOf(accessToken.getUserType());
            if (userType != UserTypeEnum.MEMBER) {
                throw new RuntimeException("不支持的用户类型: " + userType);
            }

            WebSocketUser webSocketUser = new WebSocketUser();
            webSocketUser.setUserId(accessToken.getUserId())
                    .setUserType(accessToken.getUserType())
                    .setUserInfo(accessToken.getUserInfo())
                    .setTenantId(accessToken.getTenantId())
                    .setScopes(accessToken.getScopes());

            // 构建登录用户
            return webSocketUser;
        }
        return null;
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        if (webSocketUser == null && (msg instanceof FullHttpMessage)) {
            // 握手时进行token鉴权
            try {
                Boolean checkComplete = (Boolean) ctx.channel().attr(AttributeKey.valueOf(SecurityFrameworkUtils.SECURITY_CHECK_COMPLETE_ATTRIBUTE_KEY)).get();
                if (checkComplete != null && checkComplete) {
                    super.channelRead(ctx, msg);
                    return;
                }
                // 记录连接数
                webSocketConnectionMetrics.incrementConnections();
                FullHttpMessage httpMessage = (FullHttpMessage) msg;
                HttpHeaders headers = httpMessage.headers();
                String token = Objects.requireNonNull(headers.get(SecurityFrameworkUtils.TOKEN_HEADER));
                // 开启鉴权 认证
                WebSocketUser webSocketUser = checkAccessToken(token);
                ctx.channel().attr(AttributeKey.valueOf(SecurityFrameworkUtils.WEB_SOCKET_USER_ATTRIBUTE_KEY)).set(webSocketUser);
                ctx.channel().attr(AttributeKey.valueOf(SecurityFrameworkUtils.SECURITY_CHECK_COMPLETE_ATTRIBUTE_KEY)).set(true);
                this.webSocketUser = webSocketUser;
                this.bindWebSocketSession(webSocketUser, ctx, headers);
                log.info("token 验证成功，绑定用户,userId: {}", webSocketUser.getUserId());
                super.channelRead(ctx, msg);
            } catch (Exception e) {
                refuseChannel(ctx);
                throw e;
            }
        } else {
            //other protocols
            // 发送消息时只验证channel的鉴权结果
            Boolean checkComplete = (Boolean) ctx.channel().attr(AttributeKey.valueOf(SecurityFrameworkUtils.SECURITY_CHECK_COMPLETE_ATTRIBUTE_KEY)).get();
            if (checkComplete == null || !checkComplete) {
                refuseChannel(ctx);
                return;
            }
            super.channelRead(ctx, msg);
        }
    }

    private void bindWebSocketSession(WebSocketUser user, ChannelHandlerContext ctx, HttpHeaders headers) {
        Attribute<WsSessionContext> attr = ctx.channel().attr(NettyWsChannelContext.WEB_SOCKET_USER_BIND_ATTRIBUTE_KEY);
        WsSessionContext context = new WsSessionContext();
        context.setWebSocketUser(user);
        context.setSessionId(ctx.channel().id().asLongText());
        attr.set(context);
        WebSocketAuthSuccessEvent authSuccessEvent = new WebSocketAuthSuccessEvent(this, headers);
        authSuccessEvent.setContext(context);
        authSuccessEvent.setWebSocketUser(user);

        SpringUtil.publishEvent(authSuccessEvent);

        NettyWsUserChannelCacheService userChannelCacheService = SpringUtil.getBean(NettyWsUserChannelCacheService.class);
        // 认证通过绑定用户id和当前连接
        userChannelCacheService.bindChannel(webSocketUser, ctx);

        NettyChannelCacheService nettyChannelCacheService = SpringUtil.getBean(NettyChannelCacheService.class);
        nettyChannelCacheService.addChannel(ctx, context);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        super.exceptionCaught(ctx, cause);
    }

    private void refuseChannel(ChannelHandlerContext ctx) {
        ctx.channel().writeAndFlush(new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.UNAUTHORIZED));
        ctx.channel().close();
    }

}
